#include "cnn_function.hpp"
#include "util.h"
//#include <sys/time.h>

#include <iostream>

extern std::vector<std::string> current_user_name;
extern std::vector<int> current_user_spoofing;
extern std::vector<cv::Rect> rois_global;

void uninit_network()
{
    //uninit_spoofing();
    uninit_recognition();
    uninit_detector();
    uninit_mtcnn();
}

void init_network()
{
    //init_detector("/var/darknet/cnn/cfg/yolov3-spp.cfg", "/var/darknet/cnn/yolov3-spp_final.weights");
    //init_detector("/var/darknet/cnn/cfg/yolov3-tiny.cfg", "/var/darknet/cnn/yolov3-tiny_final.weights");
    init_detector("/var/darknet/cnn/cfg/yolov3-small.cfg", "/var/darknet/cnn/model/yolov3-small_final_max_epoch_15.weights");
    init_recognition("/var/darknet/cnn/cfg/cosface_new_100.cfg", "/var/darknet/cnn/model/model.cnn.100");
    //init_recognition("/var/darknet/cnn/cfg/cosface_new.cfg", "/var/darknet/cnn/model.cnn");
    //init_spoofing("/var/darknet/cnn/cfg/spoofing.cfg", "/var/darknet/cnn/spoofing_final.weights");
    init_mtcnn("/var/darknet/cnn/cfg/mtcnn_onet.cfg", "/var/darknet/cnn/model/mtcnn_final.weights");
}

void detectLandmark_mtcnn(cv::Mat& img, std::vector<cv::Rect> &faces, std::vector<std::vector<cv::Point2f>>& extracted_landmarks)
{
    extracted_landmarks.resize(faces.size());
    int channels = img.channels();
    int face_width = 48;
    int face_height = 48;
    int face_num = faces.size();
    float *landmark = (float *)malloc(10 * face_num * sizeof(float));
    float *face_data = (float *)malloc(face_num * face_width * face_height * channels *sizeof(float));
    cv::Mat mtcnn_img;
    for(int i = 0; i < face_num; i++){
        cv::Rect mtcnn_roi = faces[i];
        //std::cout << mtcnn_roi << " " << mtcnn_roi.width << " " << mtcnn_roi.height << std::endl;
        cv::resize(img(mtcnn_roi), mtcnn_img, cv::Size(face_width, face_height));
        float *face_data_offset = face_data + i * channels * face_width * face_height;
        for(int k= 0; k < channels; ++k){
            for(int m = 0; m < face_height; ++m){
                for(int n = 0; n < face_width; ++n){
                    face_data_offset[k * face_width * face_height + n *face_width + m] =(mtcnn_img.at<cv::Vec3b>(m, n)[k] - 127.5) * 0.0078125;
                }
            }
        }
    }
    run_mtcnn(face_data, face_num, landmark);
    for(int i = 0; i < face_num; i++){
        std::vector<cv::Point2f> facial_points(5, cv::Point2f(0.f, 0.f));
        int roi_w = faces[i].width;
        int roi_h = faces[i].height;
        for(int j = 0; j < 5; j++){
            facial_points[j] = cv::Point(landmark[i * 10 + j] * roi_w + faces[i].x, landmark[i * 10 + j + 5] * roi_h +  + faces[i].y);
            //std::cout << "landmarks " << j << " " << facial_points[j] << std::endl;
            //cv::circle(img(faces[i]), facial_points[j], 4, cv::Scalar(0,0,255), -1);
        }
        extracted_landmarks[i] = facial_points;
    }
    //cv::imwrite("aaaa.jpg", img);
    free(face_data);
    free(landmark);
}

cv::Mat similarity_matric(std::vector<cv::Point2f> const& srcTri_s, std::vector<cv::Point2f> &dstTri_s)
{
    cv::Mat warp_mat(2, 3, CV_32FC1);
    int num_point = srcTri_s.size();
    //int M = srcTri_s.size();
    cv::Mat X(num_point * 2, 4, CV_32FC1);
    cv::Mat U(num_point * 2, 1, CV_32FC1);

    for (int i = 0; i < num_point; i++) {
        U.at<float>(i, 0) = srcTri_s[i].x;
        U.at<float>(i + num_point, 0) = srcTri_s[i].y;
    }

    for (int i = 0; i < num_point; i++) {
        X.at<float>(i, 0) = dstTri_s[i].x;
        X.at<float>(i, 1) = dstTri_s[i].y;
        X.at<float>(i, 2) = 1;
        X.at<float>(i, 3) = 0;

        X.at<float>(i + num_point, 0) = dstTri_s[i].y;
        X.at<float>(i + num_point, 1) = -dstTri_s[i].x;
        X.at<float>(i + num_point, 2) = 0;
        X.at<float>(i + num_point, 3) = 1;
    }

    cv::Mat X_t;
    cv::transpose(X, X_t);
    cv::Mat XX;
    XX = X_t * X;
    cv::Mat r;
    cv::Mat invXX;
    cv::invert(XX, invXX);

    r = invXX * X_t * U;
    cv::Mat Tinv(3, 3, CV_32FC1);
    float sc = r.at<float>(0, 0);
    float ss = r.at<float>(1, 0);
    float tx = r.at<float>(2, 0);
    float ty = r.at<float>(3, 0);
    Tinv.at<float>(0, 0) = sc;
    Tinv.at<float>(0, 1) = -ss;
    Tinv.at<float>(0, 2) = 0;

    Tinv.at<float>(1, 0) = ss;
    Tinv.at<float>(1, 1) = sc;
    Tinv.at<float>(1, 2) = 0;

    Tinv.at<float>(2, 0) = tx;
    Tinv.at<float>(2, 1) = ty;
    Tinv.at<float>(2, 2) = 1;

    cv::Mat T;
    cv::invert(Tinv, T);
    cv::transpose(T, T);

    cv::Mat tmp = T(cv::Rect(0, 0, 3, 2));
    return tmp;
}

int resizeInterestPoints(std::vector<cv::Point2f> const& in, std::vector<cv::Point2f> &out)
{
    double dist = cv::norm(in[0] - in[1]);
    float mTwoEyesPixel = 36;
    int zoom_in = dist / mTwoEyesPixel;
    zoom_in = zoom_in > 0 ? zoom_in : 1;
    for(size_t i = 0; i < in.size(); ++i){
        out[i].x = in[i].x / zoom_in;
        out[i].y = in[i].y / zoom_in;
    }
    return zoom_in;
}

void alignFace(std::vector<cv::Point2f> const& srcTri, cv::Mat const& frame, cv::Mat &aligned_face, cv::Rect &face_region)
{
    cv::Mat warp_mat(2, 3, CV_32FC1);
    std::vector<cv::Point2f> dstTri(5);
    dstTri[0] = cv::Point2f(30.2946f + 8, 51.6963f);
    dstTri[1] = cv::Point2f(65.5318f + 8, 51.5014f);
    dstTri[2] = cv::Point2f(48.0252f + 8, 71.7366f);
    dstTri[3] = cv::Point2f(33.5493f + 8, 92.3655f);
    dstTri[4] = cv::Point2f(62.7299f + 8, 92.2041f);
    /*
    for(int j = 0; j < 5; j++){
        cv::circle(frame(face_region), cv::Point2f(dstTri[j].x * (face_region.width / 112.0f), dstTri[j].y * (face_region.height / 112.0f)), 4, cv::Scalar(0,255,0), -1);
    }
    */
    std::vector<cv::Point2f> srcTri2(srcTri.size());
    int zoom_sz = resizeInterestPoints(srcTri, srcTri2);
    warp_mat = similarity_matric(srcTri2, dstTri);
    cv::Mat frame2;
    cv::resize(frame, frame2, cv::Size(frame.size().width / zoom_sz, frame.size().height / zoom_sz));
    //warp_mat = cv::getAffineTransform(srcTri, dstTri);
    cv::Mat warp_dst; // = cv::Mat::zeros(112, 112, frame.type());
    cv::warpAffine(frame2, warp_dst, warp_mat, cv::Size(112, 112));
    //cv::imwrite("/home/iim/face_align.jpg", warp_dst);
    aligned_face = warp_dst;
}

void detection_face(cv::Mat &img, int *total_bbox_num, int *detection_bbox)
{
    int channels = img.channels();
    cv::Mat resized;
    int network_w = 416;
    int network_h = 416;
    int new_w = img.cols;
    int new_h = img.rows;
    if (((float)network_w / img.cols) < ((float)network_h / img.rows)) {
        new_w = network_w;
        new_h = (img.rows * network_w) / img.cols;
    } else {
        new_h = network_h;
        new_w = (img.cols * network_h) / img.rows;
    }
    cv::resize(img, resized, cv::Size(new_w, new_h));
    float *image_data = (float *)malloc(resized.rows * resized.cols * channels *sizeof(float));
    for(int k= 0; k < channels; ++k){
        for(int i = 0; i < resized.rows; ++i){
            for(int j = 0; j < resized.cols; ++j){
                image_data[k*resized.rows * resized.cols + i* resized.cols + j] = resized.at<cv::Vec3b>(i, j)[channels - k - 1] / 255.0f;
            }
        }
    }
    //double start = getMsecOfNow();
    run_detection(image_data, resized.cols, resized.rows, channels, img.cols, img.rows, detection_bbox, MAX_BBOX_NUM, total_bbox_num);
    //printf("run_detection spend: %f ms\n", getMsecOfNow() - start);
    free(image_data);
}

bool rect_size_compare(cv::Rect &a, cv::Rect &b)
{
    int a_size = a.width * a.height;
    int b_size = b.width * b.height;
    if(a_size < b_size) return false;
    else return true;
}

void extract_face_feature(cv::Mat &img, int face_count, float *feature, int *detection_bbox, int *code, int is_registration)
{
    //double start = getMsecOfNow();
    std::vector<cv::Rect> faces;
    for(int i = 0; i < face_count; i++){
        cv::Rect roi_local(cv::Point(detection_bbox[i * 4], detection_bbox[i * 4 + 1]),
                           cv::Size(detection_bbox[i * 4 + 2] - detection_bbox[i * 4],
                                    detection_bbox[i * 4 + 3] - detection_bbox[i * 4 + 1]));
        faces.push_back(roi_local);
    }
    //std::sort(faces.begin(), faces.end(), rect_size_compare);
    //rois_global = faces;
    std::vector<std::vector<cv::Point2f>> extracted_landmarks;
    detectLandmark_mtcnn(img, faces, extracted_landmarks);
    std::vector<cv::Mat> aligned_faces;
    int face_width = 112;
    int face_height = 112;
    for(int i = 0; i < face_count; ++i){
        cv::Mat face_resized;
        if(is_registration == 1){
            float two_eye_height_diff = fabs((extracted_landmarks[i][1].y - extracted_landmarks[i][0].y) / faces[i].height);
            float two_lips_height_diff = fabs((extracted_landmarks[i][4].y - extracted_landmarks[i][3].y) / faces[i].height);
            float two_eye_width_diff = fabs((extracted_landmarks[i][1].x - extracted_landmarks[i][0].x) / faces[i].width);
            if(two_eye_height_diff > 0.09 ||  two_lips_height_diff > 0.9 || two_eye_width_diff < 0.30){
                printf("side face: %f %f %f\n", two_eye_height_diff, two_lips_height_diff, two_eye_width_diff);
                //*code = 4;
                //return;
            }
        }
        alignFace(extracted_landmarks[i], img, face_resized, faces[i]);
        /*
        struct timeval tv;
        gettimeofday(&tv, NULL);
        std::string time_in_micros = std::to_string(1000000 * tv.tv_sec + tv.tv_usec);
        cv::Mat combine_img, tmp_img;
        cv::resize(img(faces[i]), tmp_img, cv::Size(face_width, face_height));
        cv::hconcat(tmp_img, face_resized, combine_img);
        cv::imwrite("/home/iim/face_" + time_in_micros + ".jpg", combine_img);
        */
        //cv::imwrite("/home/iim/face_" + time_in_micros + ".jpg", img(faces[i]));
        //cv::imwrite("/home/iim/face_align_" + time_in_micros + ".jpg", face_resized);
        //std::cout << faces[i] << std::endl;
        //cv::resize(img(faces[i]), face_resized, cv::Size(face_width, face_height));
        aligned_faces.push_back(face_resized);
    }

    
    int channels = img.channels();
    float *face_data = (float *)malloc(face_count * face_width * face_height * channels *sizeof(float));
    for(int i = 0; i < face_count; i++){
        float *face_data_offset = face_data + i * face_width * face_height * channels;
        for(int k= 0; k < channels; ++k){
            for(int m = 0; m < face_height; ++m){
                for(int n = 0; n < face_width; ++n){
                    face_data_offset[k* face_width * face_height + m * face_width + n] =
                        (aligned_faces[i].at<cv::Vec3b>(m, n)[channels - k - 1] - 127.5) * 0.0078125;
                }
            }
        }
    }
    run_recognition(face_data, face_count, feature);
    /*
    for(int i = 0; i < face_count; ++i){
        cv::rectangle(img, cvPoint(faces[i].x, faces[i].y),
                      cvPoint(faces[i].x + faces[i].width, faces[i].y + faces[i].height), cvScalar(255,0,0), 4);
    }
    static int save_index = 0;
    cv::imwrite("/home/iim/testss" + std::to_string(save_index) + ".jpg", img);
    save_index += 1;
    */
    //printf("run_recognition spend: %f ms\n", getMsecOfNow() - start);
    free(face_data);
}

void check_spoofing_face(cv::Mat &img)
{
    int channels = img.channels();
    int face_width = 112;
    int face_height = 112;
    //double start = getMsecOfNow();
    float *face_data = (float *)malloc(current_user_name.size() * face_width * face_height * channels *sizeof(float));
    int face_num = 0;
    for(unsigned int i = 0; i < current_user_name.size(); i++){
        if(current_user_name[i] == "") continue;
        face_num += 1;
        cv::Mat face_region = img(rois_global[i]);
        cv::resize(face_region, face_region, cv::Size(face_width, face_height));
        float *face_data_offset = face_data + i * face_width * face_height * channels;
        for(int k= 0; k < channels; ++k){
            for(int m = 0; m < face_height; ++m){
                for(int n = 0; n < face_width; ++n){
                    face_data_offset[k* face_width * face_height + m * face_width + n] =
                        (face_region.at<cv::Vec3b>(m, n)[channels - k - 1] - 127.5) * 0.0078125;
                }
            }
        }
    }
    int *is_real_face = (int *)calloc(face_num, sizeof(int));
    run_spoofing(face_data, face_num, is_real_face);
    //printf("run_spoofing spend: %f ms\n", getMsecOfNow() - start);
    int user_name_index = 0;
    for(unsigned int i = 0; i < current_user_name.size(); i++){
        if(current_user_name[i] == "") continue;
        current_user_spoofing[i] = is_real_face[user_name_index];
        printf("spoofing is real face : %d\n", current_user_spoofing[i]);
        user_name_index += 1;
    }
    free(is_real_face);
    free(face_data);
}
